from PIL import Image
import numpy as np
import torch
from torchvision import datasets, transforms
def mask_img_3(img_path,boxes,box_id):

    img_pil=Image.open(img_path)

    img_pil_arr = np.array(img_pil)
    img_tensor = torch.tensor(img_pil_arr)
    img_for_mask = img_tensor.permute(2,0,1)
    zeros = torch.zeros_like(img_for_mask)
    for id in box_id:
        if id < len(boxes):
            zeros[:,boxes[id][0]:boxes[id][2],boxes[id][1]:boxes[id][3]]=1
    img_for_mask.masked_fill_(torch.ByteTensor(zeros),value=torch.tensor(0))
    inverse_img = img_for_mask.permute(1,2,0)
    masked_img = inverse_img.numpy()
    return masked_img

def mask_img(img_path,boxes,box_id):

    img_pil=Image.open(img_path)
    img_pil = img_pil.convert("RGB")
    h, w = img_pil.size[0], img_pil.size[1]
    img_pil = img_pil.resize((256, 256), Image.BILINEAR)
    img_pil_arr = np.array(img_pil)
    img_tensor = torch.tensor(img_pil_arr)
    # img_for_mask = img_tensor.permute(2,0,1)
    zeros = torch.zeros_like(img_tensor)
    for id in box_id:
        if id < len(boxes):
            zeros[int(boxes[id][0]/h):int(boxes[id][2]/h),int(boxes[id][1]/w):int(boxes[id][3]/w)]=1
    img_tensor.masked_fill_(torch.ByteTensor(zeros),value=torch.tensor(0))
    img_tensor = img_tensor.float()
    masked_img = img_tensor.numpy()
    return masked_img

if __name__=="__main__":

    img_path = "/mnt/myproject/pretrain/my-lm/data/funsd_data/testing_data/images/82092117.png"
    img_pil = Image.open(img_path)
    img_pil = img_pil.convert("RGB")
    print(img_pil.size)
    # h, w,_ = img_pil.size[0], img_pil.size[1],img_pil.size[2]
    img_pil = img_pil.resize((256,256),Image.BILINEAR)
    # img_pil = np.array(img_pil)

    img_array = np.array(img_pil)
    img_tensor = torch.tensor(img_array)
    img_tensor=img_tensor.float()
    print(img_tensor.numpy())
    # img_pil = img_pil.resize((100,100), Image.BILINEAR)
    # img_pil_arr = np.array(img_pil)
    #
    # img_tensor = torch.tensor(img_pil_arr)
    # zeros = torch.zeros_like(img_tensor)
    # box =[267,138,293,158]
    # zeros[int(box[0]/w):int(box[2]/w),int(box[1]/h):int(box[3]/h)]=1
    # img_tensor.masked_fill_(torch.ByteTensor(zeros), value=torch.tensor(0))
    # masked_img = img_tensor.numpy()

